import mindspore
from mindspore import nn, Tensor,ops

import numpy as np

from src.models.cswin.cswin import CSWin_64_24322_small_224
from mindspore import context

context.set_context(device_target="Ascend")
model = CSWin_64_24322_small_224()

n_parameters = sum(ops.Size()(p) for p in model.get_parameters() if p.requires_grad)

print("parameters: ", n_parameters)


skip_list = ()
if hasattr(model, 'no_weight_decay'):
    skip_list = model.no_weight_decay()
    print("model_no_weight_decay: ",skip_list)

for x in model.trainable_params():
    parameter_name = x.name
    if "token" in parameter_name:
        print("no weight decay: ", parameter_name)


x = Tensor(np.ones([1, 3, 224,224]), mindspore.float32)
y = model(x)
print("x.shape: ",y.shape)

